import os
import subprocess
import asyncio
from typing import Dict, Optional
from fastapi import FastAPI, Request, Form, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

import rospy
import rosgraph
import roslib.message  # 用于动态加载消息类型

app = FastAPI()

# 以当前文件为基准，确保无论从哪个工作目录启动，都能正确找到模板与静态资源
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
templates = Jinja2Templates(directory=os.path.join(BASE_DIR, "templates"))
app.mount("/static", StaticFiles(directory=os.path.join(BASE_DIR, "static")), name="static")

node_process: Optional[subprocess.Popen] = None
publishers: Dict[str, rospy.Publisher] = {}
subscribers: Dict[str, rospy.Subscriber] = {}
websocket_connections: Dict[str, WebSocket] = {}

@app.on_event("startup")
async def startup_event():
    if not rospy.core.is_initialized():
        rospy.init_node("fastapi_manager", anonymous=True, disable_signals=True)

@app.get("/", response_class=HTMLResponse)
async def index(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.post("/start_node")
async def start_node(file_path: str = Form(...)):
    global node_process
    # 如果先前进程对象存在但已退出，则清理引用
    if node_process is not None and node_process.poll() is not None:
        node_process = None

    if node_process is not None and node_process.poll() is None:
        return JSONResponse({"status": "error", "message": "已有节点在运行"})

    if not os.path.exists(file_path):
        return JSONResponse({"status": "error", "message": "文件不存在"})

    try:
        # 在 Windows 下无需可执行位；在类 Unix 下尽量检查可执行权限
        if os.name != "nt" and not os.access(file_path, os.X_OK):
            return JSONResponse({"status": "error", "message": "文件不可执行，请检查权限 (chmod +x)"})

        node_process = subprocess.Popen([file_path])
        return JSONResponse({"status": "ok", "message": f"已启动节点 {file_path}"})
    except Exception as exc:
        return JSONResponse({"status": "error", "message": f"启动失败: {exc}"})

@app.post("/stop_node")
async def stop_node():
    global node_process
    if node_process is None:
        return {"status": "ok", "message": "当前没有运行中的节点"}
    try:
        if node_process.poll() is None:
            # 尝试优雅结束，再强制结束
            node_process.terminate()
            try:
                node_process.wait(timeout=3)
            except Exception:
                node_process.kill()
        node_process = None
        return {"status": "ok", "message": "节点已停止"}
    except Exception as exc:
        return {"status": "error", "message": f"停止失败: {exc}"}

@app.get("/topics")
async def list_topics():
    master = rosgraph.Master("/rostopic")
    try:
        topics = master.getTopicTypes()  # [(topic, type), ...]
        return {"topics": topics}
    except:
        return {"topics": []}

@app.post("/publish")
async def publish_message(topic: str = Form(...), msg: str = Form(...)):
    global publishers
    master = rosgraph.Master("/rostopic")
    topics = dict(master.getTopicTypes())
    if topic not in topics:
        return {"status": "error", "message": f"话题 {topic} 不存在"}

    msg_type = topics[topic]
    msg_class = roslib.message.get_message_class(msg_type)
    if msg_class is None:
        return {"status": "error", "message": f"无法加载消息类型 {msg_type}"}

    if topic not in publishers:
        publishers[topic] = rospy.Publisher(topic, msg_class, queue_size=10)

    # 简单处理：只支持字符串转基本字段
    try:
        msg_obj = msg_class()
        if hasattr(msg_obj, "data"):
            msg_obj.data = msg
        else:
            return {"status": "error", "message": f"消息类型 {msg_type} 暂不支持自动填充"}
        publishers[topic].publish(msg_obj)
    except Exception as e:
        return {"status": "error", "message": str(e)}

    return {"status": "ok", "message": f"已发布到 {topic}: {msg}"}

@app.websocket("/ws/{topic}")
async def websocket_endpoint(websocket: WebSocket, topic: str):
    await websocket.accept()
    websocket_connections[topic] = websocket

    master = rosgraph.Master("/rostopic")
    topics = dict(master.getTopicTypes())
    if topic not in topics:
        await websocket.send_text(f"错误: 话题 {topic} 不存在")
        await websocket.close()
        return

    msg_type = topics[topic]
    msg_class = roslib.message.get_message_class(msg_type)
    if msg_class is None:
        await websocket.send_text(f"错误: 无法加载消息类型 {msg_type}")
        await websocket.close()
        return

    loop = asyncio.get_event_loop()

    def callback(msg):
        try:
            text = str(msg)  # 打印消息对象
            asyncio.run_coroutine_threadsafe(
                websocket.send_text(text),
                loop
            )
        except RuntimeError:
            pass

    if topic not in subscribers:
        subscribers[topic] = rospy.Subscriber(topic, msg_class, callback)

    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        websocket_connections.pop(topic, None)
        # 该实现按 topic 仅维持单连接，断开后移除订阅
        subscriber = subscribers.pop(topic, None)
        if subscriber is not None:
            try:
                subscriber.unregister()
            except Exception:
                pass

@app.get("/health")
async def health_check():
    return {"status": "ok"}

@app.on_event("shutdown")
async def shutdown_event():
    # 关闭订阅
    for topic, subscriber in list(subscribers.items()):
        try:
            subscriber.unregister()
        except Exception:
            pass
        subscribers.pop(topic, None)

    # 关闭 websocket 连接
    for topic, ws in list(websocket_connections.items()):
        try:
            await ws.close()
        except Exception:
            pass
        websocket_connections.pop(topic, None)

    # 停止可能运行的外部节点
    global node_process
    if node_process is not None and node_process.poll() is None:
        try:
            node_process.terminate()
            try:
                node_process.wait(timeout=2)
            except Exception:
                node_process.kill()
        except Exception:
            pass
        finally:
            node_process = None